#!/usr/bin/env python
# coding=utf-8
"""
安居客数据导入工具（适配版）
功能：
1. 可选择导入到MySQL或MongoDB
2. 完整支持Anjuke数据结构
3. 确保数据100%提交
"""

import os
import sys
import pymysql
from tqdm import tqdm

# 检查并导入pymongo
try:
    from pymongo import MongoClient

    MONGO_SUPPORT = True
except ImportError:
    MONGO_SUPPORT = False
    print("警告: pymongo 库未安装，MongoDB功能不可用")


def get_user_choice(options, prompt):
    """获取用户选择"""
    print(prompt)
    for i, option in enumerate(options, 1):
        print(f"{i}. {option}")
    while True:
        try:
            choice = int(input("请选择(输入数字): "))
            if 1 <= choice <= len(options):
                return choice - 1
            print("输入无效，请重新选择")
        except ValueError:
            print("请输入数字")


def connect_mysql():
    """连接MySQL数据库（适配Anjuke）"""
    try:
        conn = pymysql.connect(
            host='127.0.0.1',
            port=3306,
            user='root',
            password='123456',
            db='anjuke',
            charset='utf8mb4',
            autocommit=True
        )

        # 创建Anjuke表结构
        with conn.cursor() as cursor:
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS anjuke (
                id INT AUTO_INCREMENT PRIMARY KEY,
                date VARCHAR(20),
                district VARCHAR(50),
                area VARCHAR(50),
                xiaoqu VARCHAR(100),
                price DECIMAL(10,2),
                average_price DECIMAL(10,2),
                room_type VARCHAR(50),
                size DECIMAL(10,2),
                direction VARCHAR(50),
                decoration VARCHAR(50),
                floor VARCHAR(50),
                building_type VARCHAR(50),
                pic_url VARCHAR(255),
                detail_url VARCHAR(255),
                address TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
            """)
        return conn
    except Exception as e:
        print(f"MySQL连接失败: {e}")
        raise


def connect_mongodb():
    """连接MongoDB数据库（适配Anjuke）"""
    if not MONGO_SUPPORT:
        raise ImportError("pymongo库未安装")

    try:
        client = MongoClient('localhost', 27017)
        db = client.anjuke
        collection = db.houses

        # 创建索引
        collection.create_index([("xiaoqu", 1), ("date", 1)], unique=True)
        collection.create_index("detail_url")
        return collection
    except Exception as e:
        print(f"MongoDB连接失败: {e}")
        raise


def parse_house_info(info_str):
    """解析房屋详细信息（适配Anjuke）"""
    # Anjuke数据格式不同，这里做适配
    if not info_str:
        return {
            'room_type': '未知',
            'size': 0,
            'direction': '未知',
            'decoration': '未知',
            'floor': '未知',
            'building_type': '未知'
        }

    # 尝试从ting字段解析几室几厅
    if '室' in info_str and '厅' in info_str:
        room_type = info_str
    else:
        room_type = '未知'

    return {
        'room_type': room_type,
        'size': 0,  # 会在clean_data中单独处理
        'direction': '未知',  # 会在clean_data中单独处理
        'decoration': '未知',
        'floor': '未知',
        'building_type': '未知'  # 会在clean_data中处理
    }


def clean_data(line):
    """清洗Anjuke CSV行数据"""
    try:
        fields = line.strip().split(',')
        if len(fields) < 10:  # 确保有足够字段
            return None

        # 提取基础信息
        data = {
            'date': fields[0],  # 假设第一列是日期
            'district': '未知',  # 可以从address中提取，这里简化处理
            'area': fields[6] if len(fields) > 6 else '未知',  # area字段位置
            'xiaoqu': fields[3] if len(fields) > 3 else '',  # name字段
            'price': float(fields[4]) if len(fields) > 4 and fields[4] else 0,  # price
            'average': float(fields[5]) if len(fields) > 5 and fields[5] else 0,  # average
            'pic_url': fields[1] if len(fields) > 1 else '',  # img_url
            'detail_url': fields[2] if len(fields) > 2 else '',  # detail_url
            'address': fields[8] if len(fields) > 8 else '',  # address
            'direction': fields[7] if len(fields) > 7 else '未知',  # position
            'year': fields[9] if len(fields) > 9 else '未知'  # year
        }

        # 解析几室几厅
        ting = fields[10] if len(fields) > 10 else ''
        room_type = ting if ting else '未知'

        # 解析面积
        try:
            size = float(data['area']) if data['area'] and data['area'] != '未知' else 0
        except ValueError:
            size = 0

        # 构造房屋信息
        house_info = {
            'room_type': room_type,
            'size': size,
            'direction': data['direction'],
            'decoration': '未知',
            'floor': '未知',
            'building_type': f"{data['year']}年建" if data['year'] and data['year'] != '未知' else '未知'
        }

        return {**data, **house_info}
    except Exception as e:
        print(f"数据解析失败: {e}\n问题行: {line.strip()}")
        return None


def import_to_mysql(conn, data):
    """批量导入数据到MySQL（适配Anjuke）"""
    if not data:
        return False

    try:
        with conn.cursor() as cursor:
            sql = """
            INSERT INTO anjuke 
            (date, district, area, xiaoqu, price, average_price, room_type, 
             size, direction, decoration, floor, building_type, pic_url, detail_url, address) 
            VALUES(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            params = [
                (d.get('date', ''),
                 d.get('district', '未知'),
                 d.get('area', '未知'),
                 d.get('xiaoqu', ''),
                 d.get('price', 0),
                 d.get('average', 0),
                 d.get('room_type', '未知'),
                 d.get('size', 0),
                 d.get('direction', '未知'),
                 d.get('decoration', '未知'),
                 d.get('floor', '未知'),
                 d.get('building_type', '未知'),
                 d.get('pic_url', ''),
                 d.get('detail_url', ''),
                 d.get('address', ''))
                for d in data
            ]
            cursor.executemany(sql, params)
        conn.commit()
        print(f"成功插入 {len(data)} 条记录到MySQL")
        return True
    except Exception as e:
        print(f"MySQL插入失败: {e}")
        return False


def import_to_mongodb(collection, data):
    """批量导入数据到MongoDB（适配Anjuke）"""
    if not MONGO_SUPPORT:
        print("错误: pymongo库未安装")
        return False

    if not data:
        return False

    try:
        # 转换数据格式
        for item in data:
            item['date'] = str(item.get('date', ''))
            # 添加MongoDB需要的额外字段
            item['source'] = 'anjuke'
            item['import_time'] = datetime.datetime.now()

        result = collection.insert_many(data)
        print(f"成功插入 {len(result.inserted_ids)} 条记录到MongoDB")
        return True
    except Exception as e:
        print(f"MongoDB插入失败: {e}")
        return False


def prepare_database(db_type, db_conn, overwrite):
    """准备数据库"""
    if overwrite:
        if db_type == "mysql":
            with db_conn.cursor() as cursor:
                cursor.execute("TRUNCATE TABLE anjuke")
            print("已清空MySQL数据表")
        elif MONGO_SUPPORT:
            db_conn.delete_many({})
            print("已清空MongoDB集合")


def process_csv(csv_path, db_type, db_conn, batch_size=100):
    """处理CSV文件并导入数据"""
    if not os.path.exists(csv_path):
        print(f"错误: CSV文件不存在于 {csv_path}")
        return 0, 0

    total_lines = sum(1 for _ in open(csv_path, 'r', encoding='utf-8'))
    success = 0
    batch = []

    with open(csv_path, 'r', encoding='utf-8') as f, \
            tqdm(total=total_lines, desc="导入进度", unit="行") as pbar:

        for line in f:
            cleaned = clean_data(line)
            pbar.update(1)

            if cleaned:
                batch.append(cleaned)
                if len(batch) >= batch_size:
                    try:
                        if db_type == "mysql":
                            if import_to_mysql(db_conn, batch):
                                success += len(batch)
                        elif MONGO_SUPPORT:
                            if import_to_mongodb(db_conn, batch):
                                success += len(batch)
                        batch = []
                    except Exception as e:
                        print(f"批量插入失败: {e}")
                        batch = []

        if batch:
            try:
                if db_type == "mysql":
                    if import_to_mysql(db_conn, batch):
                        success += len(batch)
                elif MONGO_SUPPORT:
                    if import_to_mongodb(db_conn, batch):
                        success += len(batch)
            except Exception as e:
                print(f"最后批次插入失败: {e}")

    return success, total_lines


def main(csv_path):
    try:
        # 检查CSV文件
        if not os.path.exists(csv_path):
            print(f"错误: CSV文件不存在于 {csv_path}")
            return

        # 数据库类型选择
        db_options = ["MySQL"]
        if MONGO_SUPPORT:
            db_options.append("MongoDB")

        db_choice = get_user_choice(db_options, "请选择目标数据库:")
        db_type = db_options[db_choice].lower()

        # 导入方式选择
        import_modes = ["覆盖更新(先清空再导入)", "追加数据(保留原有数据)"]
        mode_choice = get_user_choice(import_modes, "请选择导入方式:")
        overwrite = (mode_choice == 0)

        # 连接数据库
        print(f"\n连接{db_options[db_choice]}数据库中...")
        if db_type == "mysql":
            db_conn = connect_mysql()
        else:
            if not MONGO_SUPPORT:
                print("错误: MongoDB支持不可用")
                return
            db_conn = connect_mongodb()

        # 准备数据库
        print("\n准备数据库中...")
        prepare_database(db_type, db_conn, overwrite)

        # 开始导入
        print(f"\n开始导入数据到{db_options[db_choice]}...")
        success, total = process_csv(csv_path, db_type, db_conn)

        # 显示结果
        print("\n导入结果:")
        print(f"总行数: {total}")
        print(f"成功导入: {success}")
        print(f"失败/跳过: {total - success}")

        # 验证数据
        if db_type == "mysql":
            with db_conn.cursor() as cursor:
                cursor.execute("SELECT COUNT(*) FROM anjuke")
                count = cursor.fetchone()[0]
                print(f"MySQL表实际记录数: {count}")
        elif MONGO_SUPPORT:
            count = db_conn.count_documents({})
            print(f"MongoDB集合实际记录数: {count}")

    except Exception as e:
        print(f"\n发生错误: {e}")
    finally:
        if 'db_conn' in locals():
            if db_type == "mysql":
                db_conn.close()
            print("数据库连接已关闭")


if __name__ == '__main__':
    import datetime  # 为MongoDB添加时间戳

    # 设置Anjuke CSV文件路径
    csv_path = "fangjia.csv"  # 修改为你的实际路径

    # 验证路径
    if not os.path.exists(csv_path):
        print(f"错误: CSV文件不存在于 {csv_path}")
        sys.exit(1)

    pymysql.install_as_MySQLdb()
    main(csv_path)